package proxy

import (
	"bytes"
	_ "embed"
	"strings"
	"sync"

	mysqlproxy "gitee.com/xuender/mysql-proxy"
	"gitee.com/xuender/mysql-proxy/pb"
	"github.com/XiaoMi/soar/database"
	"github.com/pelletier/go-toml"
	"github.com/xelabs/go-mysqlstack/driver"
	"github.com/xelabs/go-mysqlstack/sqlparser/depends/query"
	"github.com/xelabs/go-mysqlstack/sqlparser/depends/sqltypes"
	"github.com/xuender/oils/base"
	"github.com/xuender/oils/cryptos"
	"github.com/xuender/oils/logs"
)

//go:embed pass.toml
var passBytes []byte

type DBService struct {
	BaseHandler
	cons  map[uint32]*pb.Conn
	users map[uint32]*pb.User
	c     *pb.Config
	pass  *pb.Pass
	m     sync.Mutex
	cs    *CheckService
}

func NewDBService(
	config *pb.Config,
	checkService *CheckService,
) *DBService {
	pass := &pb.Pass{}
	decoder := toml.NewDecoder(bytes.NewReader(passBytes))
	base.Panic(decoder.Decode(pass))

	return &DBService{
		c:     config,
		cons:  map[uint32]*pb.Conn{},
		users: map[uint32]*pb.User{},
		pass:  pass,
		cs:    checkService,
	}
}

func (p *DBService) SQLTrim(sql string) (string, bool) {
	isPass := false
	if strings.Contains(sql, "!!") {
		isPass = true
	}

	upper := ToUpperTrum(sql)

	if !isPass {
		for _, pass := range p.pass.Contains {
			if strings.Contains(upper, pass) {
				isPass = true

				break
			}
		}
	}

	sql = database.RemoveSQLComments(sql)
	upper = ToUpperTrum(sql)

	if !isPass {
		for _, pass := range p.pass.Prefix {
			if strings.HasPrefix(upper, pass) {
				isPass = true

				break
			}
		}
	}

	return sql, isPass
}

func (p *DBService) Examine(session *driver.Session, sql string) *sqltypes.Result {
	if IsExec(sql) {
		if strings.Contains(sql, "!!") {
			return Examine()
		}

		return &sqltypes.Result{RowsAffected: 0}
	}

	return nil
}

// ComQuery 查询.
func (p *DBService) ComQuery(
	session *driver.Session,
	sql string,
	bindVariables map[string]*query.BindVariable,
	callback func(*sqltypes.Result) error,
) error {
	user, userHas := p.users[session.ID()]
	if !userHas {
		return ErrPassword
	}

	con, conHas := p.cons[session.ID()]
	if !conHas {
		return ErrPassword
	}

	newSQL, isPass := p.SQLTrim(sql)

	if err := p.Log(isPass, newSQL, con.Id, user); err != nil {
		return err
	}

	if ret := p.Check(isPass, newSQL, con.Id); ret != nil {
		return callback(ret)
	}

	if ret := p.Examine(session, newSQL); ret != nil {
		return callback(ret)
	}

	password := con.Password
	if cryptos.IsEncrypt(password) {
		password = base.Panic1(cryptos.Decrypt(password, mysqlproxy.DefaultKey))
	}

	conn := base.Panic1(driver.NewConn(
		con.Username,
		password,
		con.Address,
		con.Schema,
		con.Charset,
	))
	defer conn.Close()

	res, err := conn.FetchAll(newSQL, int(p.c.DefaultLimit))
	if err != nil {
		logs.Debugw("error", "id", session.ID(), "sql", newSQL, "err", err, "ip", session.Addr(), "closed", conn.Closed())

		return err
	}

	return callback(res)
}

func (p *DBService) Log(isPass bool, sql string, connID uint32, user *pb.User) error {
	if isPass || !user.Log {
		return nil
	}

	return p.cs.Log(sql, connID)
}

func (p *DBService) Check(isPass bool, sql string, connID uint32) *sqltypes.Result {
	if isPass {
		return nil
	}

	return p.cs.Check(sql, connID)
}

func (p *DBService) AuthCheck(session *driver.Session) error {
	p.m.Lock()
	defer p.m.Unlock()

	logs.Debugw("AuthCheck",
		"id", session.ID(),
		"User", session.User(),
		"Salt", session.Salt(),
		"Addr", session.Addr(),
		"schema", session.Schema(),
		"scramble", session.Scramble(),
	)

	for _, conn := range p.c.Conns {
		for _, user := range conn.Users {
			if user.User == session.User() && bytes.Equal(NativePassword(user.Password, session.Salt()), session.Scramble()) {
				p.users[session.ID()] = user
				p.cons[session.ID()] = conn

				logs.Debugw("AuthCheck", "ok", true)

				return nil
			}
		}
	}

	logs.Debugw("AuthCheck", "ok", false)

	return ErrPassword
}

func (p *DBService) SessionClosed(session *driver.Session) {
	p.m.Lock()
	defer p.m.Unlock()

	logs.Debugw("SessionClosed", "id", session.ID())
	delete(p.users, session.ID())
	delete(p.cons, session.ID())
}
